#!/usr/bin/env python3
"""
Lacey Index Optimization Runner
Easy-to-use script for optimizing cylinder position based on Lacey mixing index
"""

import os
import sys
import subprocess

def check_dependencies():
    """Check if required packages are available"""
    try:
        import numpy as np
        import torch
        import matplotlib.pyplot as plt
        print("✓ All dependencies available")
        return True
    except ImportError as e:
        print(f"✗ Missing dependency: {e}")
        return False

def run_optimization(target_lacey=0.6, target_timestep=100, iterations=30):
    """Run Lacey index optimization"""
    
    print("Lacey Index Cylinder Position Optimization")
    print("="*60)
    print(f"Target Lacey Index: {target_lacey}")
    print(f"Target Timestep: {target_timestep}")
    print(f"Iterations: {iterations}")
    print("="*60)
    
    # Set up paths
    sys.path.append('.')
    sys.path.append('inverse_barrier')
    
    try:
        # Import the optimizer
        from inverse_barrier.lacey_cylinder_optimizer import LaceyCylinderOptimizer
        
        # Create optimizer
        config_path = "inverse_barrier/config_lacey.json"
        optimizer = LaceyCylinderOptimizer(config_path)
        
        # Run optimization
        results = optimizer.optimize_cylinder_position(
            target_lacey=target_lacey,
            target_timestep=target_timestep,
            n_iterations=iterations
        )
        
        # Save results
        optimizer.save_results(results)
        
        # Print summary
        import numpy as np
        best_idx = np.argmin(results['losses'])
        
        print("\n" + "="*60)
        print("OPTIMIZATION COMPLETE")
        print("="*60)
        print(f"Target Lacey Index: {results['target_lacey']:.4f}")
        print(f"Best Achieved: {results['lacey_indices'][best_idx]:.4f}")
        print(f"Error: {abs(results['lacey_indices'][best_idx] - results['target_lacey']):.4f}")
        print(f"Best Angle: {results['angles'][best_idx]*180/np.pi:.1f}° ({results['angles'][best_idx]:.3f} rad)")
        print(f"Final Loss: {results['losses'][best_idx]:.6f}")
        print(f"Results: {optimizer.output_dir}")
        print("="*60)
        
        return True
        
    except Exception as e:
        print(f"Optimization failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def main():
    """Main function with user-friendly interface"""
    
    print("Lacey Index Optimization System")
    print("="*50)
    
    # Check dependencies
    if not check_dependencies():
        print("Please install missing dependencies and try again")
        return
    
    # Check model file
    if not os.path.exists("model_drum_good.pt"):
        print("✗ Model file 'model_drum_good.pt' not found")
        print("Please ensure the model file is in the root directory")
        return
    else:
        print("✓ Model file found")
    
    # Create directories
    os.makedirs("inverse_barrier/data/outputs_lacey", exist_ok=True)
    
    # Get user input or use defaults
    try:
        print("\nOptimization Parameters:")
        target_lacey = float(input("Target Lacey Index (0.0-1.0) [default: 0.6]: ") or "0.6")
        target_timestep = int(input("Target Timestep [default: 100]: ") or "100")
        iterations = int(input("Number of iterations [default: 30]: ") or "30")
        
        # Validate inputs
        if not (0.0 <= target_lacey <= 1.0):
            print("Warning: Target Lacey index should be between 0.0 and 1.0")
        
        if target_timestep < 10:
            print("Warning: Target timestep should be at least 10")
        
        if iterations < 5:
            print("Warning: Iterations should be at least 5 for meaningful results")
        
    except (ValueError, KeyboardInterrupt):
        print("\nUsing default parameters...")
        target_lacey = 0.6
        target_timestep = 100
        iterations = 30
    
    # Run optimization
    success = run_optimization(target_lacey, target_timestep, iterations)
    
    if success:
        print("\n✓ Optimization completed successfully!")
        print("Check the outputs_lacey directory for results and visualizations.")
    else:
        print("\n✗ Optimization failed. Check the error messages above.")

if __name__ == "__main__":
    main() 